#!/usr/bin/env python
# -*- conding: utf-8 -*-

"""
@Time     : 2024/10/30 6:35
@Author   : liujingmao
@File     : 1.子图实现多智能体.py
"""

from typing import TypedDict, Any, Annotated

import dotenv
from langchain_community.tools import GoogleSerperRun
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.runnables import RunnableConfig
from langchain_openai import ChatOpenAI
from langgraph.graph import MessagesState, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition

dotenv.load_dotenv()

llm = ChatOpenAI(model="gpt-4o-mini")


class GoogleSerperArgsSchema(BaseModel):
    query: str = Field(description="执行谷歌搜索的查询语句")


google_serper = GoogleSerperRun(
    api_wrapper=GoogleSerperAPIWrapper(),
    args_schema=GoogleSerperArgsSchema,
)


def reduce_str(left: str | None, right: str | None) -> str:
    if right is not None and right != "":
        return right
    return left


class AgentState(TypedDict):
    query: Annotated[str, reduce_str]  # 原始问题
    live_content: Annotated[str, reduce_str]  # 直播文案
    xhs_content: Annotated[str, reduce_str]  # 小红书文案


class LiveAgentState(AgentState, MessagesState):
    """直播文案智能体状态"""
    pass


class XHSAgentState(AgentState):
    """小红书文案智能体状态"""
    pass


def chatbot_live(state: LiveAgentState, config: RunnableConfig) -> Any:
    """直播文案智能体聊天机器人节点"""
    # 1.创建提示模板+链应用
    prompt = ChatPromptTemplate.from_messages([
        (
            "system",
            "你是一个拥有10年经验的直播文案专家，请根据用户提供的产品整理一篇直播带货脚本文案，如果在你的知识库内找不到关于该产品的信息，可以使用搜索工具。"
        ),
        ("human", "{query}"),
        ("placeholder", "{chat_history}"),
    ])
    chain = prompt | llm.bind_tools([google_serper])

    # 2.调用链并生成ai消息
    ai_message = chain.invoke({"query": state["query"], "chat_history": state["messages"]})

    return {
        "messages": [ai_message],
        "live_content": ai_message.content,
    }


# 1.创建子图1并添加节点、添加边
live_agent_graph = StateGraph(LiveAgentState)

live_agent_graph.add_node("chatbot_live", chatbot_live)
live_agent_graph.add_node("tools", ToolNode([google_serper]))

live_agent_graph.set_entry_point("chatbot_live")
live_agent_graph.add_conditional_edges("chatbot_live", tools_condition)
live_agent_graph.add_edge("tools", "chatbot_live")


def chatbot_xhs(state: XHSAgentState, config: RunnableConfig) -> Any:
    """小红书文案智能体聊天节点"""
    # 1.创建提示模板+链
    prompt = ChatPromptTemplate.from_messages([
        ("system",
         "你是一个小红书文案大师，请根据用户传递的商品名，生成一篇关于该商品的小红书笔记文案，注意风格活泼，多使用emoji表情。"),
        ("human", "{query}"),
    ])
    chain = prompt | llm | StrOutputParser()

    # 2.调用链并生成内容更新状态
    return {"xhs_content": chain.invoke({"query": state["query"]})}


# 2.创建子图2并添加节点、添加边
xhs_agent_graph = StateGraph(XHSAgentState)

xhs_agent_graph.add_node("chatbot_xhs", chatbot_xhs)

xhs_agent_graph.set_entry_point("chatbot_xhs")
xhs_agent_graph.set_finish_point("chatbot_xhs")


# 3.创建入口图并添加节点、边
def parallel_node(state: AgentState, config: RunnableConfig) -> Any:
    return state


agent_graph = StateGraph(AgentState)
agent_graph.add_node("parallel_node", parallel_node)
agent_graph.add_node("live_agent", live_agent_graph.compile())
agent_graph.add_node("xhs_agent", xhs_agent_graph.compile())

agent_graph.set_entry_point("parallel_node")
agent_graph.add_edge("parallel_node", "live_agent")
agent_graph.add_edge("parallel_node", "xhs_agent")

agent_graph.set_finish_point("live_agent")
agent_graph.set_finish_point("xhs_agent")

# 4.编译入口图
agent = agent_graph.compile()

# 5.执行入口图并打印结果
print(agent.invoke({"query": "潮汕牛肉丸"}))

'''
{'query': '潮汕牛肉丸', 'live_content': '好的，
下面是一篇关于潮汕牛肉丸的直播带货脚本文案，适合在直播平台上使用：\n\n---\n\n🎉
【潮汕牛肉丸直播盛典】🎉\n\n👋 大家好，欢迎来到我们的直播间！
今天我们要为大家带来的是来自广东潮汕地区的绝对美味——潮汕牛肉丸！🍖\n\n🌟 
**产品介绍** 🌟\n潮汕牛肉丸，作为潮汕地区最著名的传统小吃之一，
已有近百年的历史！它的制作工艺传承自客家的牛肉丸，经过改良，形成了独特的口感和风味。
我们的牛肉丸分为两种：嫩滑的牛肉丸和有嚼劲的牛筋丸，适合各种口味的你！
😋\n\n✨ **产品特点** ✨\n1. **新鲜牛肉**：采用优质牛后腿肉，肉质细腻，
口感鲜嫩！\n2. **手工制作**：每一颗牛肉丸都是经过大厨手工捶打而成，保证了肉的弹性与鲜美！
\n3. **爆汁美味**：煮熟后，牛肉丸内部鲜汁满溢，让你每一口都充满惊喜！\n
4. **营养丰富**：富含蛋白质和多种营养素，适合各个年龄段的朋友们享用！\n\n🔥 
**购买理由** 🔥\n无论是作为下饭菜，还是火锅涮煮，潮汕牛肉丸都是绝佳的选择！现在下单，
还有机会获得我们的限量优惠哦！✨\n\n🛒 **优惠活动** 🛒\n- 
凡在直播期间购买的朋友，前50名可享受买二送一的超值活动！\n- 还会有精美的潮汕风味调料包赠送，助你在家轻松做出正宗的潮汕风味！\n\n💬 **互动环节** 💬\n在直播过程中，大家可以随时在评论区提问，关于牛肉丸的制作、吃法，我们的厨师会为大家一一解答！并且，我们还会抽取幸运观众，送出神秘大礼包！🎁\n\n🚀 **赶快行动** 🚀\n喜欢潮汕牛肉丸的朋友们，快来点击下方链接购买吧！让我们一起享受这份来自潮汕的美味，让你的味蕾大开！❤️\n\n感谢大家的观看，我们下次直播再见！记得关注我们哦！👋\n\n--- \n\n希望这篇文案能够帮助你在直播中吸引更多观众和购买者！', 'xhs_content': '🌟牛肉丸控看过来！🌟今天要给大家推荐的是超经典的潮汕牛肉丸！🥢\n\n这可是我心目中的“牛肉丸之王”！👑 每一口都是满满的牛肉香气，Q弹的口感让人一吃成主顾，根本停不下来！😋✨ \n\n🔥潮汕牛肉丸的制作工艺可是相当讲究，选用新鲜的牛肉，手工打制而成，真材实料，绝对值得信赖！每一颗都是牛肉的精华，吃上一口，肉汁四溢，简直让人陶醉！💖\n\n📸今天我特地去了一家潮汕牛肉丸店，点了他们的招牌牛肉丸火锅，汤底鲜美，牛肉丸在锅里一煮，瞬间变得更加鲜嫩！搭配上香菜、葱花和自家调制的酱料，简直是完美的享受呀！🥢🍲\n\n而且，这牛肉丸还可以搭配各种食材，米粉、蔬菜，每一种组合都让人惊喜！🌈 赶紧约上小伙伴，一起去体验潮汕牛肉丸的魅力吧！💃🕺\n\n📍你们有没有尝试过潮汕牛肉丸呢？欢迎在评论区分享你们的美食经历哦！💬❤️ 
#潮汕牛肉丸 #美食推荐 #吃货日记 #Q弹美味'}
'''
